from collections import defaultdict
from typing import Any, Dict, Union

from pydantic import UUID4
from tortoise.expressions import Q
from curd.base_curd import CurdBase
from models.user_model import User, Role, Permission, UserRelation, Family
from schemas.user_schema import UserBase, RoleBase, AccessBase, UserRelationBase, FamilyBase
from core.logger import logger
from core.common import NodeMananger


class CurdUser(CurdBase[User, UserBase, UserBase]):
    def __init__(self):
        super().__init__(model=User)
        
ctrl_user = CurdUser()

class CurdUserRelation(CurdBase[UserRelation, UserRelationBase, UserRelationBase]):
    NodeFieldsMapping = {
        "id": "id",
        "label": "full_name",
        "gender": "sex",
        "img": "header_img",
        }
    
    def __init__(self):
        super().__init__(model=UserRelation)
    
    async def get_user_family_ids(self, user_obj):
        results = {}
        result_values = []
        result_count = defaultdict(set)
        if user_obj.username == "admin":
            relations = await self.model.all().prefetch_related('family', 'parent', 'child')
        else:
            relations = await self.model.filter(Q(parent_jd=user_obj.id) | Q(child_id=user_obj.id)).prefetch_related('family', 'parent', 'child')
        for relation in relations:
            if relation.family_id:
                family_obj = await relation.family
                results[family_obj.family_id] = {
                    "uuid": family_obj.family_id,
                    "name": family_obj.family_name,
                    "memberCount": 0,
                    "createTime": family_obj.create_time,
                    "role": 'member'
                }
                result_count[family_obj.family_id].add(relation.parent.id)
                result_count[family_obj.family_id].add(relation.child.id)
        for vk, values in results.items():
            values["memberCount"] = len(result_count[vk])
            result_values.append(values)
        return result_values
    
    async def get_nodes_edges(self, family_id: str):
        relations = await self.model.filter(family_id=family_id).prefetch_related('parent', 'child')
        nodes = {}
        edges = []
        spouse_mapping = {}
        
        def replace_spouse(item):
            value = {}
            if item.id in spouse_mapping:
                item.id = spouse_mapping[item.id]
            else:
                item.id = str(item.id)
            for k, v in self.NodeFieldsMapping.items():
                value[k] = getattr(item, v, "")
            if value["gender"] == 2:
                value["gender"] = "female"
            else:
                value["gender"] = "male"
            return value
        nm = NodeMananger()
        for relation in relations:
            p_obj = await relation.parent
            c_obj = await relation.child
            _p_id = str(p_obj.id)
            _c_id = str(c_obj.id)
            style = False
            if relation.relation_type == 3:
                _c_id = f"{_p_id}-spouse"
                spouse_mapping[c_obj.id] = _c_id
                style = True
            nm.add_node(_p_id, _c_id)
            nodes[_p_id] = p_obj
            nodes[_c_id] = c_obj
            if style:
                edges.append({"source": _p_id, "target": _c_id, "style": { "lineDash": [4, 4] }})
            else:
                edges.append({"source": _p_id, "target": _c_id})
        results = {"nodes": list(map(replace_spouse, nodes.values())), "edges": edges}
        results["nodes"].append({"id": "1","label": "family", "gender": "male", "img": "/static/profile/admin.png"})
        for head in nm.get_heads():
            results["edges"].append({"source": "1", "target": head})
        return results
        
ctrl_user_ralation = CurdUserRelation()


class CurdFamily(CurdBase[Family, FamilyBase, FamilyBase]):
    def __init__(self):
        super().__init__(model=Family)
    
    async def get_with_users(self, family_id: UUID4) -> Dict:
        family_obj = await self.model.get(family_id=family_id).prefetch_related("members")
        return family_obj
        
ctrl_family = CurdFamily()


class CurdRole(CurdBase[Role, RoleBase, RoleBase]):
    def __init__(self):
        super().__init__(model=Role)
        
ctrl_role = CurdRole()


class CurdPermission(CurdBase[Permission, AccessBase, AccessBase]):
    def __init__(self):
        super().__init__(model=Permission)
        
ctrl_permission = CurdPermission()